import base64
import datetime
import time

import cv2
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import requests
from matplotlib import animation
import paddlehub as hub
from PIL import Image, ImageSequence
import numpy as np
import os

# 当前文件所在路径
basepath = os.path.dirname(__file__)


def clean_back():  # img_path, out_path
    # 测试图片路径和输出路径
    test_path = '../images/'  # 按需更改路径
    output_path = 'out_img/'  # 按需更改路径

    # 待预测图片
    test_img_names = ["baby.png"]  # 按需更改文件名
    test_img_path = [test_path + img for img in test_img_names]
    img = [cv2.imread(image_path) for image_path in test_img_path]
    module = hub.Module(name="deeplabv3p_xception65_humanseg")
    input_dict = {"image": test_img_path}

    # execute predict and print the result
    print(time.time())
    results = module.segmentation(data=input_dict, output_dir=output_path)  # 本机CPU模式下，大概5秒扣一张图
    print(time.time())

    for result in results:
        print(result)
    for i in range(1):
        print(results[i]["data"].shape)
        prediction = results[i]["data"]
        plt.imshow(prediction)
        plt.show()
        newimg = np.zeros(img[i].shape)
        newimg[:, :, 0] = img[i][:, :, 0] * (prediction > 0)
        newimg[:, :, 1] = img[i][:, :, 1] * (prediction > 0)
        newimg[:, :, 2] = img[i][:, :, 2] * (prediction > 0)
        newimg = newimg.astype(np.uint8)
        print(np.max(newimg))
        print(newimg.dtype)
        # 预测结果展示
        plt.figure(figsize=(10, 10))
        plt.imshow(newimg)
        plt.axis('off')
        plt.show()
    # # 预测结果展示
    # out_img_path = 'out_img/baby.png'  # 默认输出图片的位置
    # img = mpimg.imread(out_img_path)
    # plt.figure(figsize=(10, 10))
    # plt.imshow(img)
    # plt.axis('off')
    # plt.show()


# 去除背景
def convert(upload_path):
    file_list = [upload_path]
    files = [("image", (open(item, "rb"))) for item in file_list]
    # 指定图片分割方法为deeplabv3p_xception65_humanseg并发送post请求
    # 需要先启动paddlehub一键部署！！！：
    # hub serving start - m deeplabv3p_xception65_humanseg
    # hub serving start - -config config.json
    url = "http://127.0.0.1:8866/predict/image/deeplabv3p_xception65_humanseg"
    r = requests.post(url=url, files=files)
    t = time.time()
    filename = str(t) + '.jpg'
    results = eval(r.json()["results"])
    
    for item in results:
        mypath = os.path.join(basepath, 'static/images/target', filename)
        with open(mypath, "wb") as fp:
            fp.write(base64.b64decode(item["base64"].split(',')[-1]))
            item.pop("base64")
    return filename


def get_video():
    pass


if __name__ == '__main__':
    clean_back()
